from ..data_structures.base_structures import Map, HexCell, QuadCell
import h3
from tqdm import tqdm
import rasterio

from ..quantifier.abs_quantifier import BaseQuantifier


def grid_generator(resolution:int, path:str=None, geojson=None) -> Map:
    """
    生成指定区域空的六角格网,优先使用geojson
    geojson示例:
        {
            "type": "Polygon",
            "coordinates": [[
                [min_lon, min_lat],
                [min_lon, max_lat],
                [max_lon, max_lat],
                [max_lon, min_lat],
                [min_lon, min_lat]
            ]]
        }
    """
    if path is None and geojson is None:
        # dem路径和geojson不能同时为空
        raise ValueError("请指定区域")
    elif geojson is not None:
        # 优先使用geojson
        h3_indexes = h3.polyfill_geojson(geojson, resolution)
        map:Map = Map()
        for h3_index in tqdm(h3_indexes, desc="生成六角格网:"):
            cell:HexCell = HexCell(h3_index)
            map.add_cell(cell)
        return map
    else:
        # 读取DEM数据(wgs84坐标系)
        with rasterio.open(path) as src:
            bounds = src.bounds         # 获取边界范围

        # 获取 DEM 范围的经纬度边界
        min_lat, max_lat = bounds.bottom, bounds.top
        min_lon, max_lon = bounds.left, bounds.right

        # 构造geojson
        geojson = {
            "type": "Polygon",
            "coordinates": [[
                [min_lon, min_lat],
                [min_lon, max_lat],
                [max_lon, max_lat],
                [max_lon, min_lat],
                [min_lon, min_lat]
            ]]
        }

        h3_indexes = h3.polyfill_geojson(geojson, resolution)
        map:Map = Map()
        for h3_index in tqdm(h3_indexes, desc="生成六角格网:"):
            cell:HexCell = HexCell(h3_index)
            map.add_cell(cell)
        map.map_range = geojson
        return map

def grid_quad_generator(precision:int, path:str=None, geojson=None) -> Map:
    """
    生成指定区域空的四角格网,优先使用geojson
    geojson示例:
        {
            "type": "Polygon",
            "coordinates": [[
                [min_lon, min_lat],
                [min_lon, max_lat],
                [max_lon, max_lat],
                [max_lon, min_lat],
                [min_lon, min_lat]
            ]]
        }
    """
    if path is None and geojson is None:
        # dem路径和geojson不能同时为空
        raise ValueError("请指定区域")
    elif geojson is not None:
        # 优先使用geojson
        geohash_indexes = BaseQuantifier.geohash_polyfill_geojson(geojson, precision)
        map:Map = Map()
        for geohash_index in tqdm(geohash_indexes, desc="生成四角格网:"):
            cell:QuadCell = QuadCell(geohash_index)
            map.add_cell(cell)
        return map
    else:
        # 读取DEM数据(wgs84坐标系)
        with rasterio.open(path) as src:
            bounds = src.bounds         # 获取边界范围

        # 获取 DEM 范围的经纬度边界
        min_lat, max_lat = bounds.bottom, bounds.top
        min_lon, max_lon = bounds.left, bounds.right

        # 构造geojson
        geojson = {
            "type": "Polygon",
            'coordinates': [[
                [min_lon, min_lat],
                [max_lon, min_lat],
                [max_lon, max_lat],
                [min_lon, max_lat],
                [min_lon, min_lat]
                 ]]
        }

        geohash_indexes = BaseQuantifier.geohash_polyfill_geojson(geojson, precision)
        map:Map = Map()
        for geohash_index in tqdm(geohash_indexes, desc="生成四角格网:"):
            cell:QuadCell = QuadCell(geohash_index)
            map.add_cell(cell)
        return map